{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# StructuredDataChat\n",
    "\n",
    "- Author: [hong-seongmin](https://github.com/hong-seongmin)\n",
    "- Peer Review: \n",
    "- Proofread : [Juni Lee](https://www.linkedin.com/in/ee-juni)\n",
    "- This is a part of [LangChain OpenTutorial](https://github.com/LangChain-OpenTutorial/LangChain-OpenTutorial)\n",
    "\n",
    "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/LangChain-OpenTutorial/LangChain-OpenTutorial/blob/main/14-Chains/04-Structured-Data-Chat.ipynb)[![Open in GitHub](https://img.shields.io/badge/Open%20in%20GitHub-181717?style=flat-square&logo=github&logoColor=white)](https://github.com/LangChain-OpenTutorial/LangChain-OpenTutorial/blob/main/14-Chains/04-Structured-Data-Chat.ipynb)\n",
    "\n",
    "## Overview\n",
    "\n",
    "This notebook demonstrates various data analysis and query techniques, combining traditional data processing with advanced AI-driven insights. By integrating Pandas for data handling and LangChain for LLM-powered query resolution, it explores interactive and automated approaches to data exploration.\n",
    "\n",
    "- **Data Loading and Preparation**  \n",
    "  The notebook begins by loading a Titanic dataset using pandas. The data is preprocessed and stored in a DataFrame (```df```). This section demonstrates how to read CSV files and prepare data for analysis.\n",
    "\n",
    "- **Interactive Query with LangChain Tools**  \n",
    "  Using LangChain's ```PythonAstREPLTool```, the notebook allows direct interaction with the DataFrame via Python commands. This tool provides an intuitive way to execute queries and retrieve results dynamically within the notebook environment.\n",
    "\n",
    "- **Agent-Based Query System**  \n",
    "  A key highlight of the notebook is the creation of an agent powered by LangChain and OpenAI models. This agent processes natural language queries and translates them into actionable insights using the dataset. For instance, it calculates survival rates based on specific conditions and provides context-aware responses.\n",
    "\n",
    "- **Advanced Query Scenarios**  \n",
    "  The notebook demonstrates the handling of more complex queries, such as analyzing survival rates for subsets of passengers (e.g., male passengers, children in specific classes). This is achieved by dynamically constructing and executing queries through the LLM-driven agent.\n",
    "\n",
    "- **Streamlined Query Execution**  \n",
    "  By integrating a callback system, the notebook ensures seamless real-time interaction with the LLM during query execution. This feature enhances the user experience by providing immediate feedback and results.\n",
    "\n",
    "\n",
    "### Table of Contents\n",
    "\n",
    "- [Overview](#overview)\n",
    "- [Environment Setup](#environment-setup)\n",
    "- [Data Loading](#Data-Loading)\n",
    "- [Pandas DataFrame Agent](#Pandas-DataFrame-Agent)\n",
    "- [Two or More DataFrames](#Two-or-More-DataFrames)\n",
    "\n",
    "### References\n",
    "\n",
    "- [LangChain Documentation: How to do question answering over CSVs](https://python.langchain.com/docs/how_to/sql_csv)\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Environment Setup\n",
    "\n",
    "Set up the environment. You may refer to [Environment Setup](https://wikidocs.net/257836) for more details.\n",
    "\n",
    "**[Note]**\n",
    "- ```langchain-opentutorial``` is a package that provides a set of easy-to-use environment setup, useful functions and utilities for tutorials. \n",
    "- You can checkout the [```langchain-opentutorial```](https://github.com/LangChain-OpenTutorial/langchain-opentutorial-pypi) for more details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: Ignoring invalid distribution -angchain-community (c:\\users\\user\\appdata\\local\\programs\\python\\python310\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -orch (c:\\users\\user\\appdata\\local\\programs\\python\\python310\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -rotobuf (c:\\users\\user\\appdata\\local\\programs\\python\\python310\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -treamlit (c:\\users\\user\\appdata\\local\\programs\\python\\python310\\lib\\site-packages)\n",
      "WARNING: Error parsing dependencies of torchsde: .* suffix can only be used with `==` or `!=` operators\n",
      "    numpy (>=1.19.*) ; python_version >= \"3.7\"\n",
      "           ~~~~~~~^\n",
      "WARNING: Ignoring invalid distribution -angchain-community (c:\\users\\user\\appdata\\local\\programs\\python\\python310\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -orch (c:\\users\\user\\appdata\\local\\programs\\python\\python310\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -rotobuf (c:\\users\\user\\appdata\\local\\programs\\python\\python310\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -treamlit (c:\\users\\user\\appdata\\local\\programs\\python\\python310\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -angchain-community (c:\\users\\user\\appdata\\local\\programs\\python\\python310\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -orch (c:\\users\\user\\appdata\\local\\programs\\python\\python310\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -rotobuf (c:\\users\\user\\appdata\\local\\programs\\python\\python310\\lib\\site-packages)\n",
      "WARNING: Ignoring invalid distribution -treamlit (c:\\users\\user\\appdata\\local\\programs\\python\\python310\\lib\\site-packages)\n"
     ]
    }
   ],
   "source": [
    "%%capture --no-stderr\n",
    "%pip install langchain-opentutorial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install required packages\n",
    "from langchain_opentutorial import package\n",
    "\n",
    "package.install(\n",
    "    [\n",
    "        \"langchain\",\n",
    "        \"langchain_core\",\n",
    "        \"langchain_openai\",\n",
    "        \"langchain_experimental\",\n",
    "    ],\n",
    "    verbose=False,\n",
    "    upgrade=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Environment variables have been set successfully.\n"
     ]
    }
   ],
   "source": [
    "# Set environment variables\n",
    "from langchain_opentutorial import set_env\n",
    "\n",
    "set_env(\n",
    "    {\n",
    "        \"LANGCHAIN_TRACING_V2\": \"true\",\n",
    "        \"LANGCHAIN_ENDPOINT\": \"https://api.smith.langchain.com\",\n",
    "        \"LANGCHAIN_PROJECT\": \"04-Structured-Data-Chat\",\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Alternatively, environment variables can also be set using a ```.env``` file.\n",
    "\n",
    "**[Note]**\n",
    "\n",
    "- This is not necessary if you've already set the environment variables in the previous step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Configuration file to manage API keys as environment variables\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "# Load API key information\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>PassengerId</th>\n",
       "      <th>Survived</th>\n",
       "      <th>Pclass</th>\n",
       "      <th>Name</th>\n",
       "      <th>Sex</th>\n",
       "      <th>Age</th>\n",
       "      <th>SibSp</th>\n",
       "      <th>Parch</th>\n",
       "      <th>Ticket</th>\n",
       "      <th>Fare</th>\n",
       "      <th>Cabin</th>\n",
       "      <th>Embarked</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "      <td>Braund, Mr. Owen Harris</td>\n",
       "      <td>male</td>\n",
       "      <td>22.0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>A/5 21171</td>\n",
       "      <td>7.2500</td>\n",
       "      <td>NaN</td>\n",
       "      <td>S</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>Cumings, Mrs. John Bradley (Florence Briggs Th...</td>\n",
       "      <td>female</td>\n",
       "      <td>38.0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>PC 17599</td>\n",
       "      <td>71.2833</td>\n",
       "      <td>C85</td>\n",
       "      <td>C</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>1</td>\n",
       "      <td>3</td>\n",
       "      <td>Heikkinen, Miss. Laina</td>\n",
       "      <td>female</td>\n",
       "      <td>26.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>STON/O2. 3101282</td>\n",
       "      <td>7.9250</td>\n",
       "      <td>NaN</td>\n",
       "      <td>S</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>4</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>Futrelle, Mrs. Jacques Heath (Lily May Peel)</td>\n",
       "      <td>female</td>\n",
       "      <td>35.0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>113803</td>\n",
       "      <td>53.1000</td>\n",
       "      <td>C123</td>\n",
       "      <td>S</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>5</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "      <td>Allen, Mr. William Henry</td>\n",
       "      <td>male</td>\n",
       "      <td>35.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>373450</td>\n",
       "      <td>8.0500</td>\n",
       "      <td>NaN</td>\n",
       "      <td>S</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>886</th>\n",
       "      <td>887</td>\n",
       "      <td>0</td>\n",
       "      <td>2</td>\n",
       "      <td>Montvila, Rev. Juozas</td>\n",
       "      <td>male</td>\n",
       "      <td>27.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>211536</td>\n",
       "      <td>13.0000</td>\n",
       "      <td>NaN</td>\n",
       "      <td>S</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>887</th>\n",
       "      <td>888</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>Graham, Miss. Margaret Edith</td>\n",
       "      <td>female</td>\n",
       "      <td>19.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>112053</td>\n",
       "      <td>30.0000</td>\n",
       "      <td>B42</td>\n",
       "      <td>S</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>888</th>\n",
       "      <td>889</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "      <td>Johnston, Miss. Catherine Helen \"Carrie\"</td>\n",
       "      <td>female</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1</td>\n",
       "      <td>2</td>\n",
       "      <td>W./C. 6607</td>\n",
       "      <td>23.4500</td>\n",
       "      <td>NaN</td>\n",
       "      <td>S</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>889</th>\n",
       "      <td>890</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>Behr, Mr. Karl Howell</td>\n",
       "      <td>male</td>\n",
       "      <td>26.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>111369</td>\n",
       "      <td>30.0000</td>\n",
       "      <td>C148</td>\n",
       "      <td>C</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>890</th>\n",
       "      <td>891</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "      <td>Dooley, Mr. Patrick</td>\n",
       "      <td>male</td>\n",
       "      <td>32.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>370376</td>\n",
       "      <td>7.7500</td>\n",
       "      <td>NaN</td>\n",
       "      <td>Q</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>891 rows × 12 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "     PassengerId  Survived  Pclass  \\\n",
       "0              1         0       3   \n",
       "1              2         1       1   \n",
       "2              3         1       3   \n",
       "3              4         1       1   \n",
       "4              5         0       3   \n",
       "..           ...       ...     ...   \n",
       "886          887         0       2   \n",
       "887          888         1       1   \n",
       "888          889         0       3   \n",
       "889          890         1       1   \n",
       "890          891         0       3   \n",
       "\n",
       "                                                  Name     Sex   Age  SibSp  \\\n",
       "0                              Braund, Mr. Owen Harris    male  22.0      1   \n",
       "1    Cumings, Mrs. John Bradley (Florence Briggs Th...  female  38.0      1   \n",
       "2                               Heikkinen, Miss. Laina  female  26.0      0   \n",
       "3         Futrelle, Mrs. Jacques Heath (Lily May Peel)  female  35.0      1   \n",
       "4                             Allen, Mr. William Henry    male  35.0      0   \n",
       "..                                                 ...     ...   ...    ...   \n",
       "886                              Montvila, Rev. Juozas    male  27.0      0   \n",
       "887                       Graham, Miss. Margaret Edith  female  19.0      0   \n",
       "888           Johnston, Miss. Catherine Helen \"Carrie\"  female   NaN      1   \n",
       "889                              Behr, Mr. Karl Howell    male  26.0      0   \n",
       "890                                Dooley, Mr. Patrick    male  32.0      0   \n",
       "\n",
       "     Parch            Ticket     Fare Cabin Embarked  \n",
       "0        0         A/5 21171   7.2500   NaN        S  \n",
       "1        0          PC 17599  71.2833   C85        C  \n",
       "2        0  STON/O2. 3101282   7.9250   NaN        S  \n",
       "3        0            113803  53.1000  C123        S  \n",
       "4        0            373450   8.0500   NaN        S  \n",
       "..     ...               ...      ...   ...      ...  \n",
       "886      0            211536  13.0000   NaN        S  \n",
       "887      0            112053  30.0000   B42        S  \n",
       "888      2        W./C. 6607  23.4500   NaN        S  \n",
       "889      0            111369  30.0000  C148        C  \n",
       "890      0            370376   7.7500   NaN        Q  \n",
       "\n",
       "[891 rows x 12 columns]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_experimental.tools import PythonAstREPLTool\n",
    "\n",
    "df = pd.read_csv(\"./data/titanic.csv\")\n",
    "# Read data from the titanic.csv file and store it in a DataFrame.\n",
    "tool = PythonAstREPLTool(locals={\"df\": df})\n",
    "# Use PythonAstREPLTool to create an environment containing the local variable 'df'.\n",
    "tool.invoke(\"df\")\n",
    "# Calculate the mean value of the 'Fare' column in the 'df' DataFrame."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "hello = \"\"\"\n",
    "print(\"Hello, world!\")\n",
    "\n",
    "def add(a, b):\n",
    "    return a + b\n",
    "\n",
    "print(add(30, 40))\n",
    "\n",
    "import pandas as pd\n",
    "\n",
    "df = pd.read_csv(\"./data/titanic.csv\")\n",
    "df.head()\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hello, world!\n",
      "70\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>PassengerId</th>\n",
       "      <th>Survived</th>\n",
       "      <th>Pclass</th>\n",
       "      <th>Name</th>\n",
       "      <th>Sex</th>\n",
       "      <th>Age</th>\n",
       "      <th>SibSp</th>\n",
       "      <th>Parch</th>\n",
       "      <th>Ticket</th>\n",
       "      <th>Fare</th>\n",
       "      <th>Cabin</th>\n",
       "      <th>Embarked</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "      <td>Braund, Mr. Owen Harris</td>\n",
       "      <td>male</td>\n",
       "      <td>22.0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>A/5 21171</td>\n",
       "      <td>7.2500</td>\n",
       "      <td>NaN</td>\n",
       "      <td>S</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>Cumings, Mrs. John Bradley (Florence Briggs Th...</td>\n",
       "      <td>female</td>\n",
       "      <td>38.0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>PC 17599</td>\n",
       "      <td>71.2833</td>\n",
       "      <td>C85</td>\n",
       "      <td>C</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>1</td>\n",
       "      <td>3</td>\n",
       "      <td>Heikkinen, Miss. Laina</td>\n",
       "      <td>female</td>\n",
       "      <td>26.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>STON/O2. 3101282</td>\n",
       "      <td>7.9250</td>\n",
       "      <td>NaN</td>\n",
       "      <td>S</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>4</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>Futrelle, Mrs. Jacques Heath (Lily May Peel)</td>\n",
       "      <td>female</td>\n",
       "      <td>35.0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>113803</td>\n",
       "      <td>53.1000</td>\n",
       "      <td>C123</td>\n",
       "      <td>S</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>5</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "      <td>Allen, Mr. William Henry</td>\n",
       "      <td>male</td>\n",
       "      <td>35.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>373450</td>\n",
       "      <td>8.0500</td>\n",
       "      <td>NaN</td>\n",
       "      <td>S</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   PassengerId  Survived  Pclass  \\\n",
       "0            1         0       3   \n",
       "1            2         1       1   \n",
       "2            3         1       3   \n",
       "3            4         1       1   \n",
       "4            5         0       3   \n",
       "\n",
       "                                                Name     Sex   Age  SibSp  \\\n",
       "0                            Braund, Mr. Owen Harris    male  22.0      1   \n",
       "1  Cumings, Mrs. John Bradley (Florence Briggs Th...  female  38.0      1   \n",
       "2                             Heikkinen, Miss. Laina  female  26.0      0   \n",
       "3       Futrelle, Mrs. Jacques Heath (Lily May Peel)  female  35.0      1   \n",
       "4                           Allen, Mr. William Henry    male  35.0      0   \n",
       "\n",
       "   Parch            Ticket     Fare Cabin Embarked  \n",
       "0      0         A/5 21171   7.2500   NaN        S  \n",
       "1      0          PC 17599  71.2833   C85        C  \n",
       "2      0  STON/O2. 3101282   7.9250   NaN        S  \n",
       "3      0            113803  53.1000  C123        S  \n",
       "4      0            373450   8.0500   NaN        S  "
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tool = PythonAstREPLTool(locals={\"df\": df})\n",
    "# Use PythonAstREPLTool to create an environment containing the local variable 'df'.\n",
    "tool.invoke(hello)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Loading\n",
    "\n",
    "The ```pd.read_csv``` function is used to read the CSV file and store its contents in a structured tabular format. By calling ```df.head```, we can preview the first few rows of the dataset to understand its structure and ensure the data has been loaded correctly.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>PassengerId</th>\n",
       "      <th>Survived</th>\n",
       "      <th>Pclass</th>\n",
       "      <th>Name</th>\n",
       "      <th>Sex</th>\n",
       "      <th>Age</th>\n",
       "      <th>SibSp</th>\n",
       "      <th>Parch</th>\n",
       "      <th>Ticket</th>\n",
       "      <th>Fare</th>\n",
       "      <th>Cabin</th>\n",
       "      <th>Embarked</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "      <td>Braund, Mr. Owen Harris</td>\n",
       "      <td>male</td>\n",
       "      <td>22.0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>A/5 21171</td>\n",
       "      <td>7.2500</td>\n",
       "      <td>NaN</td>\n",
       "      <td>S</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>Cumings, Mrs. John Bradley (Florence Briggs Th...</td>\n",
       "      <td>female</td>\n",
       "      <td>38.0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>PC 17599</td>\n",
       "      <td>71.2833</td>\n",
       "      <td>C85</td>\n",
       "      <td>C</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>1</td>\n",
       "      <td>3</td>\n",
       "      <td>Heikkinen, Miss. Laina</td>\n",
       "      <td>female</td>\n",
       "      <td>26.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>STON/O2. 3101282</td>\n",
       "      <td>7.9250</td>\n",
       "      <td>NaN</td>\n",
       "      <td>S</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>4</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>Futrelle, Mrs. Jacques Heath (Lily May Peel)</td>\n",
       "      <td>female</td>\n",
       "      <td>35.0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>113803</td>\n",
       "      <td>53.1000</td>\n",
       "      <td>C123</td>\n",
       "      <td>S</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>5</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "      <td>Allen, Mr. William Henry</td>\n",
       "      <td>male</td>\n",
       "      <td>35.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>373450</td>\n",
       "      <td>8.0500</td>\n",
       "      <td>NaN</td>\n",
       "      <td>S</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   PassengerId  Survived  Pclass  \\\n",
       "0            1         0       3   \n",
       "1            2         1       1   \n",
       "2            3         1       3   \n",
       "3            4         1       1   \n",
       "4            5         0       3   \n",
       "\n",
       "                                                Name     Sex   Age  SibSp  \\\n",
       "0                            Braund, Mr. Owen Harris    male  22.0      1   \n",
       "1  Cumings, Mrs. John Bradley (Florence Briggs Th...  female  38.0      1   \n",
       "2                             Heikkinen, Miss. Laina  female  26.0      0   \n",
       "3       Futrelle, Mrs. Jacques Heath (Lily May Peel)  female  35.0      1   \n",
       "4                           Allen, Mr. William Henry    male  35.0      0   \n",
       "\n",
       "   Parch            Ticket     Fare Cabin Embarked  \n",
       "0      0         A/5 21171   7.2500   NaN        S  \n",
       "1      0          PC 17599  71.2833   C85        C  \n",
       "2      0  STON/O2. 3101282   7.9250   NaN        S  \n",
       "3      0            113803  53.1000  C123        S  \n",
       "4      0            373450   8.0500   NaN        S  "
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "\n",
    "df = pd.read_csv(\"data/titanic.csv\")\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pandas DataFrame Agent\n",
    "\n",
    "The **Pandas DataFrame Agent** allows you to interact with tabular data in a conversational and dynamic way. Using LangChain’s integration with tools like ```ChatOpenAI``` and ```create_pandas_dataframe_agent```, you can perform complex queries on your DataFrame with natural language instructions. Below, we create an agent capable of analyzing the Titanic dataset and answering questions like survival rates based on specific criteria or data summary statistics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents.agent_types import AgentType\n",
    "from langchain_experimental.agents.agent_toolkits import create_pandas_dataframe_agent\n",
    "from langchain_openai import ChatOpenAI\n",
    "from langchain.agents.agent_types import AgentType\n",
    "from langchain.callbacks.base import BaseCallbackHandler\n",
    "\n",
    "class StreamCallback(BaseCallbackHandler):\n",
    "    def on_llm_new_token(self, token: str, **kwargs):\n",
    "        print(token, end=\"\", flush=True)\n",
    "\n",
    "# # Create an agent\n",
    "agent = create_pandas_dataframe_agent(\n",
    "    ChatOpenAI(\n",
    "        temperature=0,\n",
    "        model_name=\"gpt-4o-mini\",\n",
    "        streaming=True,\n",
    "        callbacks=[StreamCallback()],\n",
    "    ),  # Model definition\n",
    "    df,  # DataFrame\n",
    "    verbose=True,  # Print reasoning steps\n",
    "    agent_type=AgentType.OPENAI_FUNCTIONS,\n",
    "    allow_dangerous_code=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `python_repl_ast` with `{'query': 'df.shape'}`\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3m(891, 12)\u001b[0mThe dataframe has 891 rows and 12 columns.\u001b[32;1m\u001b[1;3mThe dataframe has 891 rows and 12 columns.\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'input': 'What is the number of rows and columns in the data?',\n",
       " 'output': 'The dataframe has 891 rows and 12 columns.'}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Query\n",
    "agent.invoke({\"input\": \"What is the number of rows and columns in the data?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `python_repl_ast` with `{'query': \"import pandas as pd\\n\\n# Assuming df1 is the dataframe containing the Titanic data\\n# Calculate the survival rate for male passengers\\nmale_passengers = df1[df1['Sex'] == 'male']\\nsurvived_males = male_passengers['Survived'].sum()\\ntotal_males = male_passengers.shape[0]\\nsurvival_rate_male = (survived_males / total_males) * 100\\nsurvival_rate_male\"}`\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3m0.0\u001b[0mThe survival rate of male passengers is 0.0%.\u001b[32;1m\u001b[1;3mThe survival rate of male passengers is 0.0%.\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'input': 'What is the survival rate of male passengers? Provide it as a percentage.',\n",
       " 'output': 'The survival rate of male passengers is 0.0%.'}"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Query\n",
    "agent.invoke(\"What is the survival rate of male passengers? Provide it as a percentage.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `python_repl_ast` with `{'query': \"import pandas as pd\\n\\n# Sample data for df1 and df2\\ndata1 = {\\n    'PassengerId': [1, 2, 3, 4, 5],\\n    'Survived': [0, 1, 1, 1, 0],\\n    'Pclass': [3, 1, 3, 1, 3],\\n    'Name': ['Braund, Mr. Owen Harris', 'Cumings, Mrs. John Bradley (Florence Briggs Thayer)', 'Heikkinen, Miss. Laina', 'Futrelle, Mrs. Jacques Heath (Lily May Peel)', 'Allen, Mr. William Henry'],\\n    'Sex': ['male', 'female', 'female', 'female', 'male'],\\n    'Age': [22, 38, 26, 35, 35],\\n    'SibSp': [1, 1, 0, 1, 0],\\n    'Parch': [0, 0, 0, 0, 0],\\n    'Ticket': ['A/5 21171', 'PC 17599', 'STON/O2. 3101282', '113803', '373450'],\\n    'Fare': [7.25, 71.2833, 7.925, 53.1, 8.05],\\n    'Cabin': [None, 'C85', None, 'C123', None],\\n    'Embarked': ['S', 'C', 'S', 'S', 'S']\\n}\\n\\ndata2 = {\\n    'PassengerId': [1, 2, 3, 4, 5],\\n    'Survived': [0, 1, 1, 1, 0],\\n    'Pclass': [3, 1, 3, 1, 3],\\n    'Name': ['Braund, Mr. Owen Harris', 'Cumings, Mrs. John Bradley (Florence Briggs Thayer)', 'Heikkinen, Miss. Laina', 'Futrelle, Mrs. Jacques Heath (Lily May Peel)', 'Allen, Mr. William Henry'],\\n    'Sex': ['male', 'female', 'female', 'female', 'male'],\\n    'Age': [22, 38, 26, 35, 35],\\n    'SibSp': [1, 1, 0, 1, 0],\\n    'Parch': [0, 0, 0, 0, 0],\\n    'Ticket': ['A/5 21171', 'PC 17599', 'STON/O2. 3101282', '113803', '373450'],\\n    'Fare': [7.25, 71.2833, 7.925, 53.1, 8.05],\\n    'Cabin': [0, 'C85', 0, 'C123', 0],\\n    'Embarked': ['S', 'C', 'S', 'S', 'S']\\n}\\n\\n# Creating DataFrames\\n# df1 has NaN in Cabin, df2 has 0 in Cabin\\ndf1 = pd.DataFrame(data1)\\ndf2 = pd.DataFrame(data2)\\n\\n# Combine the two DataFrames\\ncombined_df = pd.concat([df1, df2])\\n\\n# Filter for male passengers under 15 in 1st or 2nd class\\nfiltered_df = combined_df[(combined_df['Sex'] == 'male') & (combined_df['Age'] < 15) & (combined_df['Pclass'].isin([1, 2]))]\\n\\n# Calculate survival rate\\nif not filtered_df.empty:\\n    survival_rate = filtered_df['Survived'].mean() * 100\\nelse:\\n    survival_rate = 0.0\\n\\nsurvival_rate\"}`\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3m0.0\u001b[0mThe survival rate of male passengers under the age of 15 who were in 1st or 2nd class is 0.0%.\u001b[32;1m\u001b[1;3mThe survival rate of male passengers under the age of 15 who were in 1st or 2nd class is 0.0%.\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'input': 'What is the survival rate of male passengers under the age of 15 who were in 1st or 2nd class? Provide it as a percentage.',\n",
       " 'output': 'The survival rate of male passengers under the age of 15 who were in 1st or 2nd class is 0.0%.'}"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Query\n",
    "agent.invoke(\n",
    "    \"What is the survival rate of male passengers under the age of 15 who were in 1st or 2nd class? Provide it as a percentage.\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `python_repl_ast` with `{'query': \"import pandas as pd\\n\\n# Sample data for df1 and df2\\n# df1 = pd.DataFrame(...)  # Assuming df1 is already defined\\n# df2 = pd.DataFrame(...)  # Assuming df2 is already defined\\n\\n# Combine the two dataframes for analysis\\ncombined_df = pd.concat([df1, df2])\\n\\n# Filter for female passengers aged between 20 and 30 in 1st class\\nfiltered_df = combined_df[(combined_df['Sex'] == 'female') & \\n                           (combined_df['Age'] >= 20) & \\n                           (combined_df['Age'] <= 30) & \\n                           (combined_df['Pclass'] == 1)]\\n\\n# Calculate survival rate\\nif len(filtered_df) > 0:\\n    survival_rate = filtered_df['Survived'].mean() * 100\\nelse:\\n    survival_rate = 0\\n\\nsurvival_rate\"}`\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3m0\u001b[0mThe survival rate of female passengers aged between 20 and 30 who were in 1st class is 0%. This means that none of the female passengers in that age group and class survived.\u001b[32;1m\u001b[1;3mThe survival rate of female passengers aged between 20 and 30 who were in 1st class is 0%. This means that none of the female passengers in that age group and class survived.\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'input': 'What is the survival rate of female passengers aged between 20 and 30 who were in 1st class? Provide it as a percentage.',\n",
       " 'output': 'The survival rate of female passengers aged between 20 and 30 who were in 1st class is 0%. This means that none of the female passengers in that age group and class survived.'}"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Query\n",
    "agent.invoke(\n",
    "    \"What is the survival rate of female passengers aged between 20 and 30 who were in 1st class? Provide it as a percentage.\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Two or More DataFrames\n",
    "\n",
    "You can perform LLM-based queries based on two or more DataFrames. When entering two or more DataFrames, enclose them in ```[]```.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>PassengerId</th>\n",
       "      <th>Survived</th>\n",
       "      <th>Pclass</th>\n",
       "      <th>Name</th>\n",
       "      <th>Sex</th>\n",
       "      <th>Age</th>\n",
       "      <th>SibSp</th>\n",
       "      <th>Parch</th>\n",
       "      <th>Ticket</th>\n",
       "      <th>Fare</th>\n",
       "      <th>Cabin</th>\n",
       "      <th>Embarked</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "      <td>Braund, Mr. Owen Harris</td>\n",
       "      <td>male</td>\n",
       "      <td>22.0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>A/5 21171</td>\n",
       "      <td>7.2500</td>\n",
       "      <td>0</td>\n",
       "      <td>S</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>Cumings, Mrs. John Bradley (Florence Briggs Th...</td>\n",
       "      <td>female</td>\n",
       "      <td>38.0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>PC 17599</td>\n",
       "      <td>71.2833</td>\n",
       "      <td>C85</td>\n",
       "      <td>C</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>1</td>\n",
       "      <td>3</td>\n",
       "      <td>Heikkinen, Miss. Laina</td>\n",
       "      <td>female</td>\n",
       "      <td>26.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>STON/O2. 3101282</td>\n",
       "      <td>7.9250</td>\n",
       "      <td>0</td>\n",
       "      <td>S</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>4</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>Futrelle, Mrs. Jacques Heath (Lily May Peel)</td>\n",
       "      <td>female</td>\n",
       "      <td>35.0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>113803</td>\n",
       "      <td>53.1000</td>\n",
       "      <td>C123</td>\n",
       "      <td>S</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>5</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "      <td>Allen, Mr. William Henry</td>\n",
       "      <td>male</td>\n",
       "      <td>35.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>373450</td>\n",
       "      <td>8.0500</td>\n",
       "      <td>0</td>\n",
       "      <td>S</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   PassengerId  Survived  Pclass  \\\n",
       "0            1         0       3   \n",
       "1            2         1       1   \n",
       "2            3         1       3   \n",
       "3            4         1       1   \n",
       "4            5         0       3   \n",
       "\n",
       "                                                Name     Sex   Age  SibSp  \\\n",
       "0                            Braund, Mr. Owen Harris    male  22.0      1   \n",
       "1  Cumings, Mrs. John Bradley (Florence Briggs Th...  female  38.0      1   \n",
       "2                             Heikkinen, Miss. Laina  female  26.0      0   \n",
       "3       Futrelle, Mrs. Jacques Heath (Lily May Peel)  female  35.0      1   \n",
       "4                           Allen, Mr. William Henry    male  35.0      0   \n",
       "\n",
       "   Parch            Ticket     Fare Cabin Embarked  \n",
       "0      0         A/5 21171   7.2500     0        S  \n",
       "1      0          PC 17599  71.2833   C85        C  \n",
       "2      0  STON/O2. 3101282   7.9250     0        S  \n",
       "3      0            113803  53.1000  C123        S  \n",
       "4      0            373450   8.0500     0        S  "
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create a sample DataFrame\n",
    "df1 = df.copy()\n",
    "df1 = df1.fillna(0)\n",
    "df1.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `python_repl_ast` with `{'query': \"import pandas as pd\\n\\n# Sample data for df1 and df2\\ndata1 = {\\n    'PassengerId': [1, 2, 3, 4, 5],\\n    'Survived': [0, 1, 1, 1, 0],\\n    'Pclass': [3, 1, 3, 1, 3],\\n    'Name': [\\n        'Braund, Mr. Owen Harris',\\n        'Cumings, Mrs. John Bradley (Florence Briggs Thayer)',\\n        'Heikkinen, Miss. Laina',\\n        'Futrelle, Mrs. Jacques Heath (Lily May Peel)',\\n        'Allen, Mr. William Henry'\\n    ],\\n    'Sex': ['male', 'female', 'female', 'female', 'male'],\\n    'Age': [22, 38, 26, 35, 35],\\n    'SibSp': [1, 1, 0, 1, 0],\\n    'Parch': [0, 0, 0, 0, 0],\\n    'Ticket': ['A/5 21171', 'PC 17599', 'STON/O2. 3101282', '113803', '373450'],\\n    'Fare': [7.25, 71.2833, 7.925, 53.1, 8.05],\\n    'Cabin': [None, 'C85', None, 'C123', None],\\n    'Embarked': ['S', 'C', 'S', 'S', 'S']\\n}\\n\\ndata2 = {\\n    'PassengerId': [1, 2, 3, 4, 5],\\n    'Survived': [0, 1, 1, 1, 0],\\n    'Pclass': [3, 1, 3, 1, 3],\\n    'Name': [\\n        'Braund, Mr. Owen Harris',\\n        'Cumings, Mrs. John Bradley (Florence Briggs Thayer)',\\n        'Heikkinen, Miss. Laina',\\n        'Futrelle, Mrs. Jacques Heath (Lily May Peel)',\\n        'Allen, Mr. William Henry'\\n    ],\\n    'Sex': ['male', 'female', 'female', 'female', 'male'],\\n    'Age': [22, 38, 26, 35, 35],\\n    'SibSp': [1, 1, 0, 1, 0],\\n    'Parch': [0, 0, 0, 0, 0],\\n    'Ticket': ['A/5 21171', 'PC 17599', 'STON/O2. 3101282', '113803', '373450'],\\n    'Fare': [7.25, 71.2833, 7.925, 53.1, 8.05],\\n    'Cabin': [0, 'C85', 0, 'C123', 0],\\n    'Embarked': ['S', 'C', 'S', 'S', 'S']\\n}\\n\\ndf1 = pd.DataFrame(data1)\\ndf2 = pd.DataFrame(data2)\\n\\n# Calculate average age for both dataframes\\navg_age_df1 = df1['Age'].mean()\\navg_age_df2 = df2['Age'].mean()\\n\\n# Calculate the difference in average age\\nage_difference = avg_age_df2 - avg_age_df1\\n\\n# Calculate the percentage difference\\npercentage_difference = (age_difference / avg_age_df1) * 100\\npercentage_difference\"}`\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3m0.0\u001b[0mThe difference in the average age between the two dataframes is 0.0%. This means that the average age in both dataframes is the same.\u001b[32;1m\u001b[1;3mThe difference in the average age between the two dataframes is 0.0%. This means that the average age in both dataframes is the same.\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'input': \"What is the difference in the average age from the 'Age' column? Calculate it as a percentage.\",\n",
       " 'output': 'The difference in the average age between the two dataframes is 0.0%. This means that the average age in both dataframes is the same.'}"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create an agent\n",
    "agent = create_pandas_dataframe_agent(\n",
    "    ChatOpenAI(\n",
    "        temperature=0,\n",
    "        model_name=\"gpt-4o-mini\",\n",
    "        streaming=True,\n",
    "        callbacks=[StreamCallback()],\n",
    "    ),\n",
    "    [df, df1],\n",
    "    verbose=True,\n",
    "    agent_type=AgentType.OPENAI_FUNCTIONS,\n",
    "    allow_dangerous_code=True\n",
    ")\n",
    "\n",
    "# Query\n",
    "agent.invoke({\"input\": \"What is the difference in the average age from the 'Age' column? Calculate it as a percentage.\"})"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
